# -*- coding: utf-8 -*-
'''
File:modeling_llama4.py
Time:2025/9/15 14:48
IDE:PyCharm
Author:Barranzi An
email:awc19930818@outlook.com
github:https://github.com/La0bALanG
Barranzi's Blog:私聊巴郎子要文档链接
requirement:(Please describle your requirement here) -->
'''

# coding=utf-8
# Copyright 2025 The LLAMA4 and HuggingFace Inc. team. All rights reserved.
#
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# 导入数学计算库
import math
# 导入数据类装饰器
from dataclasses import dataclass
# 导入类型提示
from typing import Callable, Optional, Union

# 导入PyTorch核心模块
import torch
import torch.nn as nn
import torch.nn.functional as F

# 导入Llama4视觉配置类
from transformers.models.llama4.configuration_llama4 import Llama4VisionConfig

# 导入各种transformers组件
from ...activations import ACT2FN  # 激活函数映射
from ...cache_utils import Cache, DynamicCache  # 缓存工具
from ...generation import GenerationMixin  # 生成混合类
from ...integrations import use_kernel_forward_from_hub  # 内核前向传播集成
from ...masking_utils import create_causal_mask, create_chunked_causal_mask  # 掩码工具
from ...modeling_flash_attention_utils import FlashAttentionKwargs  # Flash Attention参数
from ...modeling_layers import GradientCheckpointingLayer  # 梯度检查点层
from ...modeling_outputs import BaseModelOutput, BaseModelOutputWithPast, CausalLMOutputWithPast, ModelOutput  # 模型输出类
from ...modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update  # RoPE工具
from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel  # 注意力函数和预训练模型基类
from ...processing_utils import Unpack  # 解包工具
from ...utils import TransformersKwargs, auto_docstring, can_return_tuple, logging  # 工具类和日志
from ...utils.generic import check_model_inputs  # 输入检查
from .configuration_llama4 import Llama4Config, Llama4TextConfig  # Llama4配置类

# 获取日志记录器
logger = logging.get_logger(__name__)


class Llama4TextExperts(nn.Module):
    """Llama4文本专家网络类 - 实现MoE(混合专家)架构中的专家网络"""

    def __init__(self, config: Llama4TextConfig):
        super().__init__()
        # 专家数量
        self.num_experts = config.num_local_experts
        # 中间层维度
        self.intermediate_size = config.intermediate_size
        # 隐藏层维度
        self.hidden_size = config.hidden_size
        # 专家维度(等于中间层维度)
        self.expert_dim = self.intermediate_size
        # 门控和上投影权重矩阵 - 形状为(专家数量, 隐藏维度, 2*专家维度)
        self.gate_up_proj = nn.Parameter(torch.empty(self.num_experts, self.hidden_size, 2 * self.expert_dim))
        # 下投影权重矩阵 - 形状为(专家数量, 专家维度, 隐藏维度)
        self.down_proj = nn.Parameter(torch.empty((self.num_experts, self.expert_dim, self.hidden_size)))
        # 激活函数
        self.act_fn = ACT2FN[config.hidden_act]

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        """
        专家网络前向传播

        注意：这个实现假设输入已经按专家分组，在实际使用中需要预先排序
        计算量很大，通常需要分布式计算

        Args:
            hidden_states (torch.Tensor): 输入隐藏状态 - 形状为(batch_size * token_num, hidden_size)

        Returns:
            torch.Tensor: 专家网络输出
        """
        # 重塑输入张量以匹配专家数量维度 - (专家数量, token数量, 隐藏维度)
        hidden_states = hidden_states.view(self.gate_up_proj.shape[0], -1, self.hidden_size)
        # 批量矩阵乘法：计算门控和上投影 - (专家数量, token数量, 2*专家维度)
        gate_up = torch.bmm(hidden_states, self.gate_up_proj)
        # 将gate_up分成两部分：门控和上投影
        gate, up = gate_up.chunk(2, dim=-1)  # 注意：DTensor不支持chunk操作
        # 计算下一层状态：(上投影 * 激活(门控)) @ 下投影权重
        next_states = torch.bmm((up * self.act_fn(gate)), self.down_proj)
        # 重塑输出张量 - (总token数, 隐藏维度)
        next_states = next_states.view(-1, self.hidden_size)
        return next_states


class Llama4TextMLP(nn.Module):
    """Llama4文本MLP层 - 标准的前馈神经网络"""

    def __init__(self, config, intermediate_size=None):
        super().__init__()

        # 如果没有指定中间层大小，使用配置中的默认值
        if intermediate_size is None:
            intermediate_size = config.intermediate_size

        self.config = config
        # 门控投影层：隐藏维度 -> 中间层维度
        self.gate_proj = nn.Linear(config.hidden_size, intermediate_size, bias=False)
        # 上投影层：隐藏维度 -> 中间层维度
        self.up_proj = nn.Linear(config.hidden_size, intermediate_size, bias=False)
        # 下投影层：中间层维度 -> 隐藏维度
        self.down_proj = nn.Linear(intermediate_size, config.hidden_size, bias=False)
        # 激活函数
        self.activation_fn = ACT2FN[config.hidden_act]

    def forward(self, x):
        """MLP前向传播"""
        # 计算：activation(gate_proj(x)) * up_proj(x)
        down_proj = self.activation_fn(self.gate_proj(x)) * self.up_proj(x)
        # 通过下投影层得到最终输出
        return self.down_proj(down_proj)


class Llama4TextL2Norm(torch.nn.Module):
    """Llama4文本L2归一化层"""

    def __init__(self, eps: float = 1e-6):
        super().__init__()
        # 数值稳定性常数
        self.eps = eps

    def _norm(self, x):
        """L2归一化计算"""
        # 计算L2范数：x / sqrt(mean(x^2) + eps)
        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)

    def forward(self, x):
        """前向传播：先转换为float进行计算，再转回原类型"""
        return self._norm(x.float()).type_as(x)

    def extra_repr(self):
        """返回模块的额外表示信息"""
        return f"eps={self.eps}"


class Llama4TextRMSNorm(nn.Module):
    """Llama4文本RMS归一化层 - 等效于T5LayerNorm"""

    def __init__(self, hidden_size, eps=1e-5):
        super().__init__()
        # 数值稳定性常数
        self.eps = eps
        # 可学习的权重参数
        self.weight = nn.Parameter(torch.ones(hidden_size))

    def _norm(self, x):
        """RMS归一化计算"""
        # 计算RMS：x / sqrt(mean(x^2) + eps)
        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)

    def forward(self, x):
        """前向传播"""
        # 先进行RMS归一化，再乘以可学习权重
        output = self._norm(x.float()).type_as(x)
        return output * self.weight

    def extra_repr(self):
        """返回模块的额外表示信息"""
        return f"{tuple(self.weight.shape)}, eps={self.eps}"


class Llama4Router(nn.Linear):
    """Llama4路由器 - 用于MoE架构中选择专家"""

    def __init__(self, config):
        # 继承Linear层：隐藏维度 -> 专家数量，无偏置
        super().__init__(config.hidden_size, config.num_local_experts, bias=False)
        # 专家总数
        self.num_experts = config.num_local_experts
        # 每个token选择的top-k专家数量
        self.top_k = config.num_experts_per_tok

    def forward(self, hidden_states):
        """路由器前向传播"""
        # 计算路由logits
        router_logits = super().forward(hidden_states)
        # 选择top-k专家：返回最高的k个值和对应索引
        router_top_value, router_indices = torch.topk(router_logits, self.top_k, dim=1)
        # 创建稀疏分数矩阵：未选中的专家分数为-inf
        router_scores = torch.full_like(router_logits, float("-inf")).scatter_(1, router_indices, router_top_value)
        # 使用sigmoid激活函数将分数转换为概率
        router_scores = torch.nn.functional.sigmoid(router_scores.float()).to(router_scores.dtype)
        return router_scores, router_logits


@use_kernel_forward_from_hub("Llama4TextMoe")  # 使用Hub中的内核优化前向传播
class Llama4TextMoe(nn.Module):
    """Llama4文本混合专家(MoE)模块"""

    def __init__(self, config):
        super().__init__()
        # 每个token的top-k专家数量
        self.top_k = config.num_experts_per_tok
        # 隐藏维度
        self.hidden_dim = config.hidden_size
        # 专家总数
        self.num_experts = config.num_local_experts
        # 专家网络
        self.experts = Llama4TextExperts(config)
        # 路由器
        self.router = Llama4Router(config)
        # 共享专家(所有token都会经过)
        self.shared_expert = Llama4TextMLP(config)

    def forward(self, hidden_states):
        """MoE前向传播"""
        # 重塑输入：(batch*seq_len, hidden_dim)
        hidden_states = hidden_states.reshape(-1, self.hidden_dim)
        # 路由决策：计算每个token应该分配给哪些专家
        router_scores, router_logits = self.router(hidden_states)
        # 复制输入以匹配专家数量：(num_experts * batch*seq_len, hidden_dim)
        routed_in = hidden_states.repeat(router_scores.shape[1], 1)
        # 根据路由分数加权输入
        routed_in = routed_in * router_scores.reshape(-1, 1)
        # 通过专家网络处理
        routed_out = self.experts(routed_in)
        # 通过共享专家处理原始输入
        out = self.shared_expert(hidden_states)
        # 将专家输出重塑并求和，加到共享专家输出上
        out.add_(routed_out.reshape(router_scores.shape[1], -1, routed_out.shape[-1]).sum(dim=0))
        return out, router_logits


class Llama4TextRotaryEmbedding(nn.Module):
    """Llama4文本旋转位置编码(RoPE)"""

    def __init__(self, config: Llama4TextConfig, device=None):
        super().__init__()
        # 向后兼容：原来的"type"字段现在叫"rope_type"
        self.rope_type = "llama3" if config.rope_scaling is not None else "default"

        # 缓存的最大序列长度
        self.max_seq_len_cached = config.max_position_embeddings
        # 原始最大序列长度
        self.original_max_seq_len = config.max_position_embeddings

        self.config = config
        # 获取RoPE初始化函数
        self.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type]

        # 初始化逆频率和注意力缩放
        inv_freq, self.attention_scaling = self.rope_init_fn(self.config, device)
        # 注册为缓冲区(不会被保存为模型参数)
        self.register_buffer("inv_freq", inv_freq, persistent=False)
        # 保存原始逆频率
        self.original_inv_freq = self.inv_freq

    @torch.no_grad()  # 禁用梯度计算
    @dynamic_rope_update  # 动态RoPE更新装饰器(用于高级RoPE类型)
    def forward(self, x, position_ids):
        """RoPE前向传播"""
        # 扩展逆频率：(1, freq_dim, 1) -> (batch_size, freq_dim, 1)
        inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1)
        # 扩展位置ID：(batch_size, seq_len) -> (batch_size, 1, seq_len)
        position_ids_expanded = position_ids[:, None, :].float()

        # 确定设备类型(MPS设备用CPU处理)
        device_type = x.device.type if isinstance(x.device.type, str) and x.device.type != "mps" else "cpu"
        # 强制使用float32以确保数值稳定性
        with torch.autocast(device_type=device_type, enabled=False):
            # 计算频率：逆频率 @ 位置ID
            freqs = (inv_freq_expanded.to(x.device) @ position_ids_expanded).transpose(1, 2)
            # 转换为复数表示：创建单位圆上的复数
            freqs_cis = torch.polar(torch.ones_like(freqs), freqs)
            # 应用注意力缩放
            freqs_cis = freqs_cis * self.attention_scaling

        return freqs_cis


def apply_rotary_emb(
        xq: torch.Tensor,  # 查询张量
        xk: torch.Tensor,  # 键张量
        freqs_cis: torch.Tensor,  # 频率复数
) -> tuple[torch.Tensor, torch.Tensor]:
    """应用旋转位置编码到查询和键张量"""
    # 将实数张量转换为复数：(..., dim) -> (..., dim//2, 2) -> (..., dim//2)
    xq_ = torch.view_as_complex(xq.float().reshape(*xq.shape[:-1], -1, 2))
    xk_ = torch.view_as_complex(xk.float().reshape(*xk.shape[:-1], -1, 2))
    # 应用旋转：复数乘法实现旋转
    xq_out = torch.view_as_real(xq_ * freqs_cis[:, :, None, :]).flatten(3)
    xk_out = torch.view_as_real(xk_ * freqs_cis[:, :, None, :]).flatten(3)
    # 转回原始数据类型
    return xq_out.type_as(xq), xk_out.type_as(xk)


def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
    """
    重复键值张量以匹配查询头数量
    等价于torch.repeat_interleave(x, dim=1, repeats=n_rep)
    从(batch, num_key_value_heads, seqlen, head_dim)到(batch, num_attention_heads, seqlen, head_dim)
    """
    batch, num_key_value_heads, slen, head_dim = hidden_states.shape
    # 如果重复次数为1，直接返回
    if n_rep == 1:
        return hidden_states
    # 扩展维度并重复
    hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)
    # 重塑为目标形状
    return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)


def eager_attention_forward(
        module: nn.Module,  # 注意力模块
        query: torch.Tensor,  # 查询张量
        key: torch.Tensor,  # 键张量
        value: torch.Tensor,  # 值张量
        attention_mask: Optional[torch.Tensor],  # 注意力掩码
        scaling: float,  # 缩放因子
        dropout: float = 0.0,  # dropout概率
        **kwargs,
):
    """
    Eager注意力前向传播实现
    改编自transformers.models.llama.modeling_llama.eager_attention_forward
    llama4不将注意力权重转换为fp32
    """
    # 重复键值以匹配查询头数量
    key_states = repeat_kv(key, module.num_key_value_groups)
    value_states = repeat_kv(value, module.num_key_value_groups)

    # 计算注意力分数：Q @ K^T * scaling
    attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling
    # 应用因果掩码
    if attention_mask is not None:
        causal_mask = attention_mask[:, :, :, : key_states.shape[-2]]
        attn_weights = attn_weights + causal_mask

    # Softmax归一化
    attn_weights = nn.functional.softmax(attn_weights, dim=-1)
    # 应用dropout
    attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training)
    # 计算注意力输出：注意力权重 @ V
    attn_output = torch.matmul(attn_weights, value_states)
    # 调整维度顺序
    attn_output = attn_output.transpose(1, 2).contiguous()

    return attn_output, attn_weights


def vision_eager_attention_forward(
        module: nn.Module,  # 视觉注意力模块
        query: torch.Tensor,  # 查询张量
        key: torch.Tensor,  # 键张量
        value: torch.Tensor,  # 值张量
        attention_mask: Optional[torch.Tensor],  # 注意力掩码
        scaling: float,  # 缩放因子(实际未使用)
        dropout: float = 0.0,  # dropout概率
        **kwargs,
):
    """
    视觉模型的Eager注意力前向传播实现
    改编自transformers.models.llama.modeling_llama.eager_attention_forward
    llama4不将注意力权重转换为fp32
    """
    # 重复键值以匹配查询头数量
    key_states = repeat_kv(key, module.num_key_value_groups)
    value_states = repeat_kv(value, module.num_key_value_groups)

    # 计算注意力分数：使用head_dim^-0.5作为缩放因子(而非传入的scaling参数)
    attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * module.head_dim ** -0.5
    # 应用因果掩码
    if attention_mask is not None:
        causal_mask = attention_mask[:, :, :, : key_states.shape[-2]]
        attn_weights = attn_weights + causal_mask

    # Softmax归一化
    attn_weights = nn.functional.softmax(attn_weights, dim=-1)
    # 应用dropout
    attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training)
    # 计算注意力输出
    attn_output = torch.matmul(attn_weights, value_states)
    # 调整维度顺序
    attn_output = attn_output.transpose(1, 2).contiguous()

    return attn_output, attn_weights


class Llama4TextAttention(nn.Module):
    """来自'Attention Is All You Need'论文的多头注意力机制"""

    def __init__(self, config: Llama4TextConfig, layer_idx):
        super().__init__()
        # 保存配置对象
        self.config = config
        # 当前层的索引号，用于确定该层的特殊配置
        self.layer_idx = layer_idx
        # 每个注意力头的维度，可以是配置中指定的head_dim或者通过hidden_size/num_attention_heads计算得出
        self.head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads)
        # 注意力头的总数量
        self.num_attention_heads = config.num_attention_heads
        # 键值头组的数量，用于分组查询注意力(GQA)，计算方式为注意力头数除以键值头数
        self.num_key_value_groups = config.num_attention_heads // config.num_key_value_heads
        # 键值头的数量，通常小于注意力头数量以节省计算和内存
        self.num_key_value_heads = config.num_key_value_heads
        # 注意力缩放因子，使用head_dim的-0.5次方，用于缩放注意力分数
        self.scaling = self.head_dim ** -0.5
        # 注意力缩放参数，用于温度调节
        self.attn_scale = config.attn_scale
        # 底部缩放参数，用于温度调节中的计算
        self.floor_scale = config.floor_scale
        # 是否启用注意力温度调节机制
        self.attn_temperature_tuning = config.attn_temperature_tuning
        # 注意力dropout概率
        self.attention_dropout = config.attention_dropout
        # 标记这是因果注意力（单向注意力）
        self.is_causal = True
        # 根据层索引确定该层是否使用RoPE位置编码
        self.use_rope = config.no_rope_layers[layer_idx]

        # 查询(Query)投影层：将hidden_size维度映射到(num_attention_heads * head_dim)维度
        self.q_proj = nn.Linear(
            config.hidden_size, config.num_attention_heads * self.head_dim, bias=config.attention_bias
        )
        # 键(Key)投影层：将hidden_size维度映射到(num_key_value_heads * head_dim)维度
        self.k_proj = nn.Linear(
            config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias
        )
        # 值(Value)投影层：将hidden_size维度映射到(num_key_value_heads * head_dim)维度
        self.v_proj = nn.Linear(
            config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias
        )
        # 输出投影层：将(num_attention_heads * head_dim)维度映射回hidden_size维度
        self.o_proj = nn.Linear(
            config.num_attention_heads * self.head_dim, config.hidden_size, bias=config.attention_bias
        )
        # 如果配置启用QK归一化并且该层使用RoPE，则初始化QK归一化层
        if self.config.use_qk_norm and self.use_rope:
            self.qk_norm = Llama4TextL2Norm(config.rms_norm_eps)

    def forward(
            self,
            hidden_states: torch.Tensor,  # 输入的隐藏状态张量
            position_embeddings: tuple[torch.Tensor, torch.Tensor],  # 位置编码元组
            attention_mask: Optional[torch.Tensor],  # 注意力掩码，用于屏蔽某些位置
            past_key_value: Optional[Cache] = None,  # 过去的键值缓存，用于加速推理
            cache_position: Optional[torch.LongTensor] = None,  # 缓存位置索引
            **kwargs: Unpack[FlashAttentionKwargs],  # Flash Attention的额外参数
    ) -> tuple[torch.Tensor, Optional[torch.Tensor], Optional[tuple[torch.Tensor]]]:
        """注意力机制的前向传播函数"""

        # 获取输入张量的形状，不包括最后一个维度（特征维度）
        input_shape = hidden_states.shape[:-1]  # (batch_size, seq_len)
        # 构造隐藏状态的目标形状，用于重塑查询和值张量
        hidden_shape = (*input_shape, -1, self.head_dim)  # (batch_size, seq_len, num_heads, head_dim)

        # 通过查询投影层计算查询状态，并重塑为多头形式
        query_states = self.q_proj(hidden_states).view(hidden_shape)
        # 通过键投影层计算键状态，并重塑为多头形式
        key_states = self.k_proj(hidden_states).view(*input_shape, -1, self.head_dim)
        # 通过值投影层计算值状态，重塑后转置seq_len和num_heads维度
        value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)

        # 如果该层使用RoPE位置编码（16E模型在某些层跳过RoPE以处理长上下文）
        if self.use_rope:
            # 对查询和键状态应用旋转位置编码
            query_states, key_states = apply_rotary_emb(
                query_states, key_states, position_embeddings.to(query_states.device)
            )

        # 如果存在QK归一化层（128E模型不使用qk_norm）
        if hasattr(self, "qk_norm"):
            # 对查询状态进行L2归一化
            query_states = self.qk_norm(query_states)
            # 对键状态进行L2归一化
            key_states = self.qk_norm(key_states)

        # 使用温度调节机制（来自论文https://huggingface.co/papers/2501.19399），仅用于NoROPE层
        if self.attn_temperature_tuning and not self.use_rope:
            # 计算基于位置的注意力缩放因子：log(floor((pos+1)/floor_scale) + 1) * attn_scale + 1
            attn_scales = (
                    torch.log(
                        torch.floor((cache_position.float() + 1.0) / self.floor_scale) + 1.0) * self.attn_scale + 1.0
            )
            # 将缩放因子重塑并扩展以匹配查询张量的形状（支持batch_size > 1）
            attn_scales = attn_scales.view((1, input_shape[-1], 1, 1)).expand((*input_shape, 1, 1))
            # 将缩放因子应用到查询状态上
            query_states = (query_states * attn_scales).to(query_states.dtype)

        # 转置查询状态的维度：(batch, seq_len, num_heads, head_dim) -> (batch, num_heads, seq_len, head_dim)
        query_states = query_states.transpose(1, 2)
        # 转置键状态的维度：(batch, seq_len, num_heads, head_dim) -> (batch, num_heads, seq_len, head_dim)
        key_states = key_states.transpose(1, 2)

        # 如果存在过去的键值缓存
        if past_key_value is not None:
            # 对于RoPE模型需要sin和cos信息；静态缓存需要cache_position参数
            cache_kwargs = {"cache_position": cache_position}
            # 更新键值缓存，返回更新后的键状态和值状态
            key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)

        # 选择注意力实现方式，默认使用eager实现
        attention_interface: Callable = eager_attention_forward
        # 如果配置指定了其他注意力实现方式，则使用相应的实现
        if self.config._attn_implementation != "eager":
            attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]

        # 执行注意力计算
        attn_output, attn_weights = attention_interface(
            self,  # 当前注意力模块
            query_states,  # 查询状态
            key_states,  # 键状态
            value_states,  # 值状态
            attention_mask,  # 注意力掩码
            dropout=0.0 if not self.training else self.attention_dropout,  # 训练时使用dropout，推理时不使用
            scaling=self.scaling,  # 缩放因子
            **kwargs,  # 其他参数
        )

        # 将注意力输出重塑回原始形状并确保内存连续性
        attn_output = attn_output.reshape(*input_shape, -1).contiguous()
        # 通过输出投影层得到最终输出
        attn_output = self.o_proj(attn_output)
        # 返回注意力输出和注意力权重
        return attn_output, attn_weights


class Llama4TextDecoderLayer(GradientCheckpointingLayer):
    """Llama4文本解码器层，继承自梯度检查点层以支持内存优化"""

    def __init__(self, config, layer_idx):
        super().__init__()
        # 隐藏层维度
        self.hidden_size = config.hidden_size
        # 当前层索引
        self.layer_idx = layer_idx
        # 当前层的注意力类型（全注意力或分块注意力）
        self.attention_type = config.layer_types[layer_idx]
        # 自注意力模块
        self.self_attn = Llama4TextAttention(config, layer_idx)
        # 判断当前层是否为MoE层
        self.is_moe_layer = layer_idx in config.moe_layers

        # 根据是否为MoE层选择不同的前馈网络
        if self.is_moe_layer:  # 128E模型交替使用密集层和稀疏层
            self.feed_forward = Llama4TextMoe(config)
        else:
            # 使用标准MLP，使用配置中的intermediate_size_mlp作为中间层大小
            self.feed_forward = Llama4TextMLP(config, intermediate_size=config.intermediate_size_mlp)

        # 输入层归一化，应用在自注意力之前
        self.input_layernorm = Llama4TextRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        # 注意力后归一化，应用在前馈网络之前
        self.post_attention_layernorm = Llama4TextRMSNorm(config.hidden_size, eps=config.rms_norm_eps)

    def forward(
            self,
            hidden_states: torch.Tensor,  # 输入隐藏状态
            attention_mask: Optional[torch.Tensor] = None,  # 注意力掩码
            position_ids: Optional[torch.LongTensor] = None,  # 位置ID
            past_key_value: Optional[tuple[torch.Tensor]] = None,  # 过去的键值对
            use_cache: Optional[bool] = False,  # 是否使用缓存
            cache_position: Optional[torch.LongTensor] = None,  # 缓存位置
            position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]] = None,  # 位置编码，必要但保留用于向后兼容
            **kwargs: Unpack[FlashAttentionKwargs],  # Flash Attention参数
    ) -> tuple[torch.FloatTensor, Optional[tuple[torch.FloatTensor, torch.FloatTensor]]]:
        """解码器层的前向传播"""

        # 保存残差连接的输入
        residual = hidden_states

        # 应用输入层归一化
        hidden_states = self.input_layernorm(hidden_states)

        # 自注意力计算
        attention_states, _ = self.self_attn(
            hidden_states=hidden_states,  # 归一化后的隐藏状态
            position_embeddings=position_embeddings,  # 位置编码
            attention_mask=attention_mask,  # 注意力掩码
            past_key_value=past_key_value,  # 键值缓存
            use_cache=use_cache,  # 是否使用缓存
            cache_position=cache_position,  # 缓存位置
            **kwargs,  # 其他参数
        )
        # 残差连接：将注意力输出加到原始输入上
        hidden_states = residual + attention_states

        # 前馈网络部分
        # 保存新的残差连接输入
        residual = hidden_states
        # 应用注意力后归一化
        hidden_states = self.post_attention_layernorm(hidden_states)
        # 通过前馈网络
        hidden_states = self.feed_forward(hidden_states)
        # 如果是MoE层，feed_forward返回元组，需要提取第一个元素
        if self.is_moe_layer:
            hidden_states, _ = hidden_states
        # 残差连接：确保形状匹配后相加
        hidden_states = residual + hidden_states.view(residual.shape)
        # 返回最终的隐藏状态
        return hidden_states


@auto_docstring  # 自动生成文档字符串
class Llama4PreTrainedModel(PreTrainedModel):
    """Llama4预训练模型基类"""

    # 配置类型注解
    config: Llama4Config
    # 支持梯度检查点以节省内存
    supports_gradient_checkpointing = True
    # 在设备放置时跳过的键，past_key_values通常保持在原设备上
    _skip_keys_device_placement = ["past_key_values"]
    # 不支持Flash Attention
    _supports_flash_attn = False
    # 支持缩放点积注意力(SDPA)
    _supports_sdpa = True
    # 支持灵活注意力
    _supports_flex_attn = True

    # 可以编译完整图
    _can_compile_fullgraph = True
    # 支持注意力后端
    _supports_attention_backend = True

    def _init_weights(self, module):
        """初始化模型权重的函数"""
        # 获取初始化标准差，优先使用config.initializer_range，否则使用text_config中的
        std = (
            self.config.initializer_range
            if hasattr(self.config, "initializer_range")
            else self.config.text_config.initializer_range
        )

        # 线性层的权重初始化
        if isinstance(module, nn.Linear):
            # 权重使用正态分布初始化，均值为0，标准差为std
            module.weight.data.normal_(mean=0.0, std=std)
            # 如果有偏置项，初始化为0
            if module.bias is not None:
                module.bias.data.zero_()
        # 嵌入层的权重初始化
        elif isinstance(module, nn.Embedding):
            # 权重使用正态分布初始化
            module.weight.data.normal_(mean=0.0, std=std)
            # 如果有padding索引，将对应位置的权重设为0
            if module.padding_idx is not None:
                module.weight.data[module.padding_idx].zero_()
        # LayerNorm层的初始化
        elif isinstance(module, nn.LayerNorm):
            # 权重初始化为1
            module.weight.data.fill_(1.0)
            # 偏置初始化为0
            module.bias.data.zero_()
        # RMSNorm层的初始化
        elif isinstance(module, Llama4TextRMSNorm):
            # 权重初始化为1
            module.weight.data.fill_(1.0)
        # 专家网络的权重初始化
        elif isinstance(module, Llama4TextExperts):
            # 门控上投影权重使用正态分布初始化
            module.gate_up_proj.data.normal_(mean=0.0, std=std)
            # 下投影权重使用正态分布初始化
            module.down_proj.data.normal_(mean=0.0, std=std)
        # 视觉模型的特殊初始化
        elif isinstance(module, Llama4VisionModel):
            # 类嵌入使用缩放的正态分布初始化
            module.class_embedding.data.normal_(std=module.scale)
            # 位置嵌入使用缩放的正态分布初始化
            module.positional_embedding_vlm.data.normal_(std=module.scale)


@auto_docstring  # 自动生成文档字符串
class Llama4TextModel(Llama4PreTrainedModel):
    """Llama4文本模型类"""

    # 不拆分的模块列表，用于模型并行
    _no_split_modules = ["Llama4TextDecoderLayer"]
    # 基础模型前缀
    base_model_prefix = "model"
    # 配置类型注解
    config: Llama4TextConfig
    # 可记录输出的组件映射
    _can_record_outputs = {
        "attentions": Llama4TextAttention,  # 注意力输出
        "hidden_states": Llama4TextDecoderLayer,  # 隐藏状态输出
        "router_logits": Llama4TextMoe,  # 路由器logits输出
    }

    def __init__(self, config: Llama4TextConfig):
        super().__init__(config)
        # padding token的索引
        self.padding_idx = config.pad_token_id
        # 词汇表大小
        self.vocab_size = config.vocab_size

        # 词嵌入层：将token ID映射到hidden_size维度的向量
        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
        # 创建所有解码器层的模块列表
        self.layers = nn.ModuleList(
            [Llama4TextDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
        )
        # 最终的RMS归一化层
        self.norm = Llama4TextRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        # 旋转位置编码
        self.rotary_emb = Llama4TextRotaryEmbedding(config=config)
        # 梯度检查点标志
        self.gradient_checkpointing = False

        # 初始化权重并应用最终处理
        self.post_init()

    @check_model_inputs  # 检查模型输入
    @auto_docstring  # 自动生成文档字符串
    def forward(
            self,
            input_ids: torch.LongTensor = None,  # 输入的token ID
            attention_mask: Optional[torch.Tensor] = None,  # 注意力掩码
            position_ids: Optional[torch.LongTensor] = None,  # 位置ID
            past_key_values: Optional[Cache] = None,  # 过去的键值缓存
            inputs_embeds: Optional[torch.FloatTensor] = None,  # 输入嵌入（可选）
            use_cache: Optional[bool] = None,  # 是否使用缓存
            cache_position: Optional[torch.LongTensor] = None,  # 缓存位置
            **kwargs: Unpack[TransformersKwargs],  # 其他transformers参数
    ) -> Union[tuple, BaseModelOutputWithPast]:
        """文本模型的前向传播"""

        # 检查input_ids和inputs_embeds必须指定且只能指定一个
        if (input_ids is None) ^ (inputs_embeds is not None):
            raise ValueError("You must specify exactly one of input_ids or inputs_embeds")

        # 如果没有提供inputs_embeds，通过input_ids获取嵌入
        if inputs_embeds is None:
            inputs_embeds = self.embed_tokens(input_ids.to(self.embed_tokens.weight.device))

        # 如果使用缓存但没有提供past_key_values，创建动态缓存
        if use_cache and past_key_values is None:
            past_key_values = DynamicCache()

        # 如果没有提供cache_position，计算当前位置
        if cache_position is None:
            # 获取已处理的token数量
            past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
            # 创建当前序列的位置索引
            cache_position = torch.arange(
                past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device
            )

        # 如果没有提供position_ids，从cache_position创建
        if position_ids is None:
            position_ids = cache_position.unsqueeze(0)

        # 处理注意力掩码，可能已经被prepare()等函数预处理过
        if not isinstance(causal_mask_mapping := attention_mask, dict):
            # 准备掩码参数
            mask_kwargs = {
                "config": self.config,
                "input_embeds": inputs_embeds,
                "attention_mask": attention_mask,
                "cache_position": cache_position,
                "past_key_values": past_key_values,
                "position_ids": position_ids,
            }
            # 创建不同类型的因果掩码
            causal_mask_mapping = {
                "full_attention": create_causal_mask(**mask_kwargs),  # 全注意力掩码
                "chunked_attention": create_chunked_causal_mask(**mask_kwargs),  # 分块注意力掩码
            }

        # 初始化隐藏状态为输入嵌入
        hidden_states = inputs_embeds

        # 创建在所有解码器层间共享的位置编码
        freq_cis = self.rotary_emb(hidden_states, position_ids)

        # 依次通过所有解码器层
        for decoder_layer in self.layers[: self.config.num_hidden_layers]:
            hidden_states = decoder_layer(
                hidden_states,  # 当前隐藏状态
                attention_mask=causal_mask_mapping[decoder_layer.attention_type],  # 根据层类型选择掩码
                position_ids=position_ids,  # 位置ID
                past_key_value=past_key_values,  # 键值缓存
                use_cache=use_cache,  # 是否使用缓存
                cache_position=cache_position,  # 缓存位置
                position_embeddings=freq_cis,  # 位置编码
                **kwargs,  # 其他参数
            )

        # 应用最终的归一化
        hidden_states = self.norm(hidden_states)

        # 返回模型输出
        return BaseModelOutputWithPast(
            last_hidden_state=hidden_states,  # 最后的隐藏状态
            past_key_values=past_key_values if use_cache else None,  # 键值缓存（如果使用）
        )


class Llama4ForCausalLM(Llama4PreTrainedModel, GenerationMixin):
    """Llama4因果语言模型类，用于文本生成任务，继承自预训练模型基类和生成混合类"""

    # 不拆分的模块列表，用于模型并行时保持这些模块的完整性
    _no_split_modules = ["Llama4TextDecoderLayer"]
    # 基础模型前缀，用于标识底层模型
    base_model_prefix = "language_model"
    # 绑定权重的键列表，lm_head权重通常与embedding权重绑定以减少参数
    _tied_weights_keys = ["lm_head.weight"]
    # 张量并行计划，指定lm_head使用列向并行复制策略
    _tp_plan = {"lm_head": "colwise_rep"}
    # 配置类型注解
    config: Llama4TextConfig

    def __init__(self, config: Llama4TextConfig):
        """初始化因果语言模型"""
        super().__init__(config)
        # 创建底层的Llama4文本模型，负责编码和注意力计算
        self.model = Llama4TextModel(config)
        # 保存词汇表大小
        self.vocab_size = config.vocab_size
        # 语言模型头：线性层，将隐藏状态(hidden_size)映射到词汇表大小(vocab_size)，用于预测下一个token
        # bias=False表示不使用偏置项，这是Llama模型的标准设置
        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)

        # 初始化权重并应用最终处理（如权重绑定等）
        self.post_init()

    def set_decoder(self, decoder):
        """设置解码器模块"""
        self.model = decoder

    def get_decoder(self):
        """获取解码器模块"""
        return self.model

    @can_return_tuple  # 装饰器：允许函数返回元组格式
    @auto_docstring  # 装饰器：自动生成文档字符串
    def forward(
            self,
            input_ids: torch.LongTensor = None,  # 输入的token ID序列，形状为(batch_size, sequence_length)
            attention_mask: Optional[torch.Tensor] = None,  # 注意力掩码，用于指示哪些token需要注意
            position_ids: Optional[torch.LongTensor] = None,  # 位置ID，指定每个token在序列中的位置
            past_key_values: Optional[Union[Cache, list[torch.FloatTensor]]] = None,  # 过去的键值缓存，用于加速推理
            inputs_embeds: Optional[torch.FloatTensor] = None,  # 预计算的输入嵌入，可替代input_ids
            labels: Optional[torch.LongTensor] = None,  # 标签序列，用于计算语言建模损失
            use_cache: Optional[bool] = None,  # 是否使用键值缓存来加速推理
            cache_position: Optional[torch.LongTensor] = None,  # 缓存中的位置索引
            logits_to_keep: Union[int, torch.Tensor] = 0,  # 要保留的logits数量，用于内存优化
            **kwargs: Unpack[TransformersKwargs],  # 其他transformers相关参数
    ) -> Union[tuple, CausalLMOutputWithPast]:
        r"""
        labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
            用于计算掩码语言建模损失的标签。索引应该在`[0, ..., config.vocab_size]`范围内或为-100
            （参见`input_ids`文档字符串）。索引为`-100`的token会被忽略（掩码），
            损失只对标签在`[0, ..., config.vocab_size]`范围内的token计算。

        Example:

        ```python
        >>> from transformers import AutoTokenizer, Llama4ForCausalLM

        >>> model = Llama4ForCausalLM.from_pretrained("meta-llama4/Llama4-2-7b-hf")
        >>> tokenizer = AutoTokenizer.from_pretrained("meta-llama4/Llama4-2-7b-hf")

        >>> prompt = "Hey, are you conscious? Can you talk to me?"
        >>> inputs = tokenizer(prompt, return_tensors="pt")

        >>> # Generate
        >>> generate_ids = model.generate(inputs.input_ids, max_length=30)
        >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
        "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you."
        ```"""

        # 通过底层文本模型进行前向传播，获取隐藏状态
        outputs = self.model(
            input_ids=input_ids,
            attention_mask=attention_mask,
            position_ids=position_ids,
            past_key_values=past_key_values,
            inputs_embeds=inputs_embeds,
            use_cache=use_cache,
            cache_position=cache_position,
            **kwargs,
        )

        # 提取最后一层的隐藏状态，形状为(batch_size, sequence_length, hidden_size)
        hidden_states = outputs[0]

        # 只计算必要的logits以优化内存使用，如果不计算损失则不将其上转换为float
        # 如果logits_to_keep是整数，则取最后logits_to_keep个位置；否则使用提供的索引
        slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
        # 通过语言模型头计算词汇表上的logits，形状为(batch_size, kept_length, vocab_size)
        logits = self.lm_head(hidden_states[:, slice_indices, :])

        # 初始化损失为None
        loss = None
        # 如果提供了标签，计算语言建模损失
        if labels is not None:
            # 调用损失函数计算交叉熵损失
            loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.config.vocab_size, **kwargs)

        # 返回因果语言模型的输出，包含损失、logits和其他状态
        return CausalLMOutputWithPast(
            loss=loss,  # 计算得到的损失值
            logits=logits,  # 预测的logits
            past_key_values=outputs.past_key_values,  # 键值缓存
            hidden_states=outputs.hidden_states,  # 隐藏状态
            attentions=outputs.attentions,  # 注意力权重
        )


@dataclass  # 数据类装饰器，自动生成__init__、__repr__等方法
@auto_docstring(  # 自动文档字符串装饰器
    custom_intro="""
    Llava因果语言模型（或自回归）输出的基类。
    """
)
class Llama4CausalLMOutputWithPast(ModelOutput):
    """Llama4因果语言模型输出类，扩展了标准的ModelOutput以支持图像相关功能"""
    r"""
    loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided):
        语言建模损失（用于下一个token预测）。
    logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.vocab_size)`):
        语言建模头的预测分数（SoftMax之前每个词汇token的分数）。
    past_key_values (`Cache`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`):
        长度为`config.n_layers`的`tuple(torch.FloatTensor)`元组，每个元组有2个形状为
        `(batch_size, num_heads, sequence_length, embed_size_per_head)`的张量

        包含预计算的隐藏状态（自注意力块中的键和值），可用于加速顺序解码（参见`past_key_values`输入）。
    image_hidden_states (`torch.FloatTensor`, *optional*):
        大小为(batch_size, num_images, sequence_length, hidden_size)`的`torch.FloatTensor`。
        模型由视觉编码器产生并在投影最后隐藏状态后的image_hidden_states。
    """

    # 语言建模损失，可选
    loss: Optional[torch.FloatTensor] = None
    # 预测logits
    logits: torch.FloatTensor = None
    # 过去的键值缓存
    past_key_values: Optional[list[torch.FloatTensor]] = None
    # 隐藏状态
    hidden_states: Optional[tuple[torch.FloatTensor]] = None
    # 注意力权重
    attentions: Optional[tuple[torch.FloatTensor]] = None
    # 图像隐藏状态，用于多模态任务
    image_hidden_states: Optional[torch.FloatTensor] = None


class Llama4VisionMLP2(torch.nn.Module):
    """Llama4视觉MLP第二版，用于视觉特征的非线性变换"""

    def __init__(self, config):
        super().__init__()
        # 隐藏层维度大小
        self.hidden_size = config.hidden_size
        # 中间层维度大小
        self.intermediate_size = config.intermediate_size
        # 第一个全连接层：从中间维度映射到投影器输入维度
        self.fc1 = nn.Linear(self.intermediate_size, config.projector_input_dim, bias=False)
        # 第二个全连接层：从投影器输出维度映射到投影器输出维度（维度保持不变）
        self.fc2 = nn.Linear(config.projector_output_dim, config.projector_output_dim, bias=False)
        # 激活函数使用GELU（高斯误差线性单元）
        self.activation_fn = nn.GELU()  # ACT2FN[config.hidden_act]
        # dropout概率
        self.dropout = config.projector_dropout

    def forward(self, hidden_states):
        """MLP前向传播"""
        # 通过第一个全连接层
        hidden_states = self.fc1(hidden_states)
        # 应用GELU激活函数
        hidden_states = self.activation_fn(hidden_states)
        # 应用dropout进行正则化（仅在训练时）
        hidden_states = F.dropout(hidden_states, p=self.dropout, training=self.training)
        # 通过第二个全连接层并再次应用激活函数
        return self.activation_fn(self.fc2(hidden_states))


class Llama4MultiModalProjector(nn.Module):
    """Llama4多模态投影器，用于将视觉特征投影到文本空间"""

    def __init__(self, config):
        super().__init__()
        # 线性投影层：将视觉编码器的输出维度映射到文本模型的隐藏维度
        # 这是连接视觉和文本模态的关键组件
        self.linear_1 = nn.Linear(
            config.vision_config.vision_output_dim,  # 视觉编码器输出维度
            config.text_config.hidden_size,  # 文本模型隐藏维度
            bias=False,  # 不使用偏置
        )

    def forward(self, image_features):
        """投影器前向传播"""
        # 将图像特征投影到文本特征空间
        hidden_states = self.linear_1(image_features)
        return hidden_states


def pixel_shuffle(input_tensor, shuffle_ratio):
    """
    像素重排函数，用于重新排列图像patch的空间布局

    Args:
        input_tensor: 输入张量，形状为[batch_size, num_patches, channels]
        shuffle_ratio: 重排比例

    Returns:
        重排后的张量
    """
    # 获取输入张量的形状：[batch_size, num_patches, channels]
    batch_size, num_patches, channels = input_tensor.shape
    # 计算patch大小（假设patch排列为正方形）
    patch_size = int(math.sqrt(num_patches))

    # 重塑为2D patch网格：[batch_size, patch_size, patch_size, channels]
    input_tensor = input_tensor.view(batch_size, patch_size, patch_size, -1)
    batch_size, height, width, channels = input_tensor.size()

    # 第一次重塑：改变宽度和通道数的分布
    reshaped_tensor = input_tensor.view(batch_size, height, int(width * shuffle_ratio), int(channels / shuffle_ratio))
    # 转置维度：交换宽度和高度维度
    reshaped_tensor = reshaped_tensor.permute(0, 2, 1, 3).contiguous()

    # 第二次重塑：进一步调整空间和通道维度
    reshaped_tensor = reshaped_tensor.view(
        batch_size, int(height * shuffle_ratio), int(width * shuffle_ratio), int(channels / (shuffle_ratio ** 2))
    )
    # 再次转置维度
    reshaped_tensor = reshaped_tensor.permute(0, 2, 1, 3).contiguous()

    # 最终重塑为原始格式：[batch_size, new_num_patches, new_channels]
    output_tensor = reshaped_tensor.view(batch_size, -1, reshaped_tensor.shape[-1])
    return output_tensor


class Llama4VisionPixelShuffleMLP(nn.Module):
    """Llama4视觉像素重排MLP，结合像素重排和MLP处理"""

    def __init__(self, config):
        super().__init__()
        # 像素重排比例
        self.pixel_shuffle_ratio = config.pixel_shuffle_ratio
        # 内部维度：根据重排比例调整投影器输入维度
        self.inner_dim = int(config.projector_input_dim // (self.pixel_shuffle_ratio ** 2))
        # 输出维度
        self.output_dim = config.projector_output_dim
        # MLP处理模块
        self.mlp = Llama4VisionMLP2(config)

    def forward(self, encoded_patches: torch.Tensor) -> torch.Tensor:
        """前向传播：先进行像素重排，再通过MLP处理"""
        # 对编码的patch进行像素重排
        encoded_patches = pixel_shuffle(encoded_patches, self.pixel_shuffle_ratio)
        # 通过MLP进行进一步处理
        return self.mlp(encoded_patches)


# TODO: 这里有一个专门用于视觉编码器的不同RoPE实现，定义如下
def reshape_for_broadcast(freqs_ci: torch.Tensor, query: torch.Tensor):
    """
    重塑频率张量以进行广播

    Args:
        freqs_ci: 频率复数张量
        query: 查询张量

    Returns:
        重塑后的频率张量
    """
    # 获取查询张量的维度数
    ndim = query.ndim
    # 创建广播形状：只保留第1维和最后一维，其余维度设为1
    shape = [d if i == 1 or i == ndim - 1 else 1 for i, d in enumerate(query.shape)]
    # 重塑频率张量以匹配广播形状
    return freqs_ci.view(*shape)


def vision_apply_rotary_emb(
        query: torch.Tensor,  # 查询张量
        key: torch.Tensor,  # 键张量
        freqs_ci: torch.Tensor,  # 频率复数张量
) -> tuple[torch.Tensor, torch.Tensor]:
    """
    对视觉模型应用旋转位置编码

    Args:
        query: 查询张量
        key: 键张量
        freqs_ci: 频率复数张量

    Returns:
        应用RoPE后的查询和键张量
    """
    # 将查询张量转换为复数形式：(..., dim) -> (..., dim//2, 2) -> (..., dim//2)
    query_ = torch.view_as_complex(query.float().reshape(*query.shape[:-1], -1, 2))
    # 将键张量转换为复数形式
    key_ = torch.view_as_complex(key.float().reshape(*key.shape[:-1], -1, 2))
    # 重塑频率张量以进行广播：freqs_ci[:,:,None,:]
    freqs_ci = reshape_for_broadcast(freqs_ci=freqs_ci, query=query_)
    # 确保频率张量在正确的设备上
    freqs_ci = freqs_ci.to(query_.device)

    # 应用旋转：复数乘法实现旋转变换
    query_out = torch.view_as_real(query_ * freqs_ci).flatten(3)
    key_out = torch.view_as_real(key_ * freqs_ci).flatten(3)

    # 转回原始数据类型并返回
    return query_out.type_as(query), key_out.type_as(key)  # 但这会降到8e-3精度


class Llama4VisionAttention(nn.Module):
    """Llama4视觉注意力模块，专门用于处理视觉特征的多头自注意力机制"""

    def __init__(self, config: Llama4VisionConfig):
        """初始化视觉注意力模块"""
        super().__init__()
        # 保存配置对象
        self.config = config
        # 嵌入维度，等于隐藏层大小
        self.embed_dim = config.hidden_size
        # 注意力头的数量
        self.num_heads = config.num_attention_heads
        # 每个注意力头的维度，通过总维度除以头数得到
        self.head_dim = config.hidden_size // config.num_attention_heads
        # 键值组数量，视觉注意力中通常为1（标准多头注意力）
        self.num_key_value_groups = 1
        # 注意力机制中的dropout概率
        self.attention_dropout = config.attention_dropout
        # 注意力分数的缩放因子，使用head_dim的-0.5次方
        self.scaling = self.head_dim ** -0.5

        # 查询(Query)投影层：将输入维度映射到所有头的查询维度，带偏置
        self.q_proj = nn.Linear(self.embed_dim, self.num_heads * self.head_dim, bias=True)
        # 键(Key)投影层：将输入维度映射到所有头的键维度，带偏置
        self.k_proj = nn.Linear(self.embed_dim, self.num_heads * self.head_dim, bias=True)
        # 值(Value)投影层：将输入维度映射到所有头的值维度，带偏置
        self.v_proj = nn.Linear(self.embed_dim, self.num_heads * self.head_dim, bias=True)
        # 输出投影层：将所有头的输出合并回原始嵌入维度，带偏置
        self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.embed_dim, bias=True)

    def forward(
            self,
            hidden_states: torch.Tensor,  # 输入的隐藏状态，形状为(batch_size, seq_len, hidden_size)
            freqs_ci: torch.Tensor,  # 用于旋转位置编码的复数频率张量
            attention_mask: Optional[torch.Tensor] = None,  # 注意力掩码，用于屏蔽特定位置
            past_key_value: Optional[Cache] = None,  # 过去的键值缓存（视觉模型中通常不使用）
            **kwargs: Unpack[FlashAttentionKwargs],  # Flash Attention的额外参数
    ) -> tuple[torch.Tensor, Optional[torch.Tensor], Optional[tuple[torch.Tensor]]]:
        """视觉注意力模块的前向传播"""

        # 获取输入张量的形状，不包括最后一个维度（特征维度）
        # input_shape = (batch_size, seq_len)
        input_shape = hidden_states.shape[:-1]
        # 构造重塑后的形状，用于多头注意力
        # hidden_shape = (batch_size, seq_len, num_heads, head_dim)
        hidden_shape = (*input_shape, -1, self.head_dim)

        # 通过投影层计算查询、键、值状态，并重塑为多头形式
        query_states = self.q_proj(hidden_states).view(hidden_shape)
        key_states = self.k_proj(hidden_states).view(hidden_shape)
        value_states = self.v_proj(hidden_states).view(hidden_shape)

        # 对查询和键状态应用视觉专用的旋转位置编码（RoPE）
        # 这与文本模型的RoPE实现有所不同，专门针对视觉特征优化
        query_states, key_states = vision_apply_rotary_emb(query_states, key_states, freqs_ci=freqs_ci)

        # 转置维度以匹配注意力计算的要求
        # 从(batch, seq_len, num_heads, head_dim)转为(batch, num_heads, seq_len, head_dim)
        query_states = query_states.transpose(1, 2)
        key_states = key_states.transpose(1, 2)
        value_states = value_states.transpose(1, 2)

        # 选择注意力实现方式，默认使用视觉专用的eager实现
        attention_interface: Callable = vision_eager_attention_forward
        # 禁用flex_attention，因为在张量并行度为8时会出现问题
        # 原因是嵌入维度为88，不是2的幂次方
        if self.config._attn_implementation not in ["eager", "flex_attention"]:
            attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]

        # 执行注意力计算
        attn_output, attn_weights = attention_interface(
            self,  # 当前注意力模块实例
            query_states,  # 查询状态张量
            key_states,  # 键状态张量
            value_states,  # 值状态张量
            None,  # 注意力掩码（视觉注意力通常不使用掩码）
            dropout=0.0 if not self.training else self.attention_dropout,  # 训练时使用dropout
            scaling=None,  # TODO: 可能需要在此强制执行以保证张量并行兼容性，因为缩放不仅仅是sqrt(head_dim)
            is_causal=False,  # 必须强制设为False，视觉注意力不是因果的（可以看到所有位置）
            **kwargs,  # 其他Flash Attention参数
        )

        # 将注意力输出重塑回原始形状并确保内存连续性
        attn_output = attn_output.reshape(*input_shape, -1).contiguous()
        # 通过输出投影层获得最终输出
        attn_output = self.o_proj(attn_output)
        # 返回注意力输出和注意力权重
        return attn_output, attn_weights


class Llama4VisionMLP(nn.Module):
    """Llama4视觉MLP模块，用于视觉编码器中的前馈网络层"""

    def __init__(self, config):
        """初始化视觉MLP模块"""
        super().__init__()
        # 保存配置对象
        self.config = config
        # 激活函数使用GELU（高斯误差线性单元）
        self.activation_fn = nn.GELU()  # ACT2FN[config.hidden_act]
        # 第一个全连接层：从隐藏维度扩展到中间维度，带偏置
        # 这个扩展通常使隐藏维度增大4倍，提供更强的表达能力
        self.fc1 = nn.Linear(config.hidden_size, config.intermediate_size, bias=True)
        # 第二个全连接层：从中间维度压缩回隐藏维度，带偏置
        self.fc2 = nn.Linear(config.intermediate_size, config.hidden_size, bias=True)

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        """MLP前向传播：线性变换 -> 激活 -> 线性变换"""
        # 通过第一个全连接层进行维度扩展
        hidden_states = self.fc1(hidden_states)
        # 应用GELU激活函数，引入非线性
        hidden_states = self.activation_fn(hidden_states)
        # 通过第二个全连接层压缩回原始维度
        hidden_states = self.fc2(hidden_states)
        return hidden_states


class Llama4VisionEncoderLayer(GradientCheckpointingLayer):
    """Llama4视觉编码器层，继承自梯度检查点层以支持内存优化"""

    def __init__(self, config: Llama4VisionConfig):
        """初始化视觉编码器层"""
        super().__init__()
        # 隐藏层维度大小
        self.hidden_size = config.hidden_size

        # 自注意力模块，用于建模patch之间的关系
        self.self_attn = Llama4VisionAttention(config)
        # MLP模块，用于非线性变换和特征增强
        self.mlp = Llama4VisionMLP(config)

        # 输入层归一化，应用在自注意力之前
        # 注意：视觉模型使用LayerNorm而不是RMSNorm
        self.input_layernorm = nn.LayerNorm(config.hidden_size)
        # 注意力后层归一化，应用在MLP之前
        self.post_attention_layernorm = nn.LayerNorm(config.hidden_size)

    def forward(
            self,
            hidden_state: torch.Tensor,  # 输入隐藏状态，通常是图像patch的嵌入
            freqs_ci: torch.Tensor,  # 用于RoPE的频率复数张量
            attention_mask: Optional[torch.Tensor] = None,  # 注意力掩码（可选）
            output_attentions: Optional[bool] = None,  # 是否输出注意力权重
    ):
        """视觉编码器层的前向传播，采用Pre-LN架构"""

        # === 自注意力部分 ===
        # 保存残差连接的输入
        residual = hidden_state

        # 应用输入层归一化（Pre-LN架构）
        hidden_state = self.input_layernorm(hidden_state)

        # 通过自注意力模块处理归一化后的隐藏状态
        hidden_state, attn_weights = self.self_attn(
            hidden_state,  # 归一化后的隐藏状态
            freqs_ci=freqs_ci,  # RoPE频率张量
            attention_mask=attention_mask,  # 注意力掩码
        )
        # 残差连接：将注意力输出与原始输入相加
        hidden_state = residual + hidden_state

        # === 前馈网络部分 ===
        # 保存新的残差连接输入
        residual = hidden_state
        # 应用注意力后层归一化
        hidden_state = self.post_attention_layernorm(hidden_state)
        # 通过MLP进行非线性变换
        hidden_state = self.mlp(hidden_state)
        # 残差连接：将MLP输出与输入相加
        hidden_state = residual + hidden_state

        # 准备输出元组，至少包含隐藏状态
        outputs = (hidden_state,)

        # 如果需要输出注意力权重，将其添加到输出中
        if output_attentions:
            outputs += (attn_weights,)

        return outputs


class Llama4VisionEncoder(nn.Module):
    """
    Llama4视觉编码器，由多个自注意力层组成的Transformer编码器
    包含 `config.num_hidden_layers` 个自注意力层，每一层都是 [`Llama4VisionEncoderLayer`]

    Args:
        config: Llama4VisionConfig 配置对象
    """

    def __init__(self, config: Llama4VisionConfig):
        """初始化视觉编码器"""
        super().__init__()
        # 保存配置对象
        self.config = config
        # 创建编码器层的模块列表，包含指定数量的编码器层
        # 每一层都是独立的Llama4VisionEncoderLayer实例
        self.layers = nn.ModuleList([Llama4VisionEncoderLayer(config) for _ in range(config.num_hidden_layers)])
        # 是否启用梯度检查点，用于节省显存（以计算时间换内存）
        self.gradient_checkpointing = False
        # 再次保存配置（代码重复，可能是历史原因）
        self.config = config

    def forward(
            self,
            hidden_states: torch.Tensor,  # 输入的隐藏状态，形状为(batch_size, seq_len, hidden_size)
            freqs_ci: torch.Tensor,  # 用于RoPE的复数频率张量 TODO: 考虑将其作为属性而不是参数传递
            attention_mask: Optional[torch.Tensor] = None,  # 注意力掩码，可选
            output_attentions: Optional[bool] = None,  # 是否输出注意力权重
            output_hidden_states: Optional[bool] = None,  # 是否输出所有层的隐藏状态
            return_dict: Optional[bool] = None,  # 是否返回字典格式的输出
    ) -> Union[tuple, BaseModelOutput]:
        r"""
        视觉编码器的前向传播

        Args:
            inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`):
                可选，直接传入嵌入表示而不是input_ids。当你想要更好地控制如何将input_ids索引
                转换为相关向量时很有用，而不是使用模型内部的嵌入查找矩阵。
            attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
                用于避免在填充token索引上执行注意力的掩码。掩码值在 `[0, 1]` 中选择：

                - 1 表示 **未被掩码** 的token，
                - 0 表示 **被掩码** 的token。

                [什么是注意力掩码？](../glossary#attention-mask)
            output_attentions (`bool`, *optional*):
                是否返回所有注意力层的注意力张量。详见返回张量中的 `attentions`。
            output_hidden_states (`bool`, *optional*):
                是否返回所有层的隐藏状态。详见返回张量中的 `hidden_states`。
            return_dict (`bool`, *optional*):
                是否返回 [`~utils.ModelOutput`] 而不是普通元组。
        """
        # 如果未指定输出选项，则使用配置中的默认值
        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
        output_hidden_states = (
            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
        )
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict

        # 初始化存储中间状态的变量
        encoder_states = () if output_hidden_states else None  # 存储所有层的隐藏状态
        all_attentions = () if output_attentions else None  # 存储所有层的注意力权重

        # 逐层处理，通过所有编码器层
        for encoder_layer in self.layers:
            # 如果需要输出隐藏状态，保存当前层的输入状态
            if output_hidden_states:
                encoder_states = encoder_states + (hidden_states,)

            # 通过当前编码器层进行前向传播
            layer_outputs = encoder_layer(
                hidden_state=hidden_states,  # 当前隐藏状态
                attention_mask=attention_mask,  # 注意力掩码
                output_attentions=output_attentions,  # 是否输出注意力权重
                freqs_ci=freqs_ci,  # RoPE频率张量
            )

            # 如果需要输出注意力权重，保存当前层的注意力权重
            if output_attentions:
                all_attentions = all_attentions + (layer_outputs[1],)

            # 更新隐藏状态为当前层的输出
            hidden_states = layer_outputs[0]

        # 保存最后一层的隐藏状态
        if output_hidden_states:
            encoder_states = encoder_states + (hidden_states,)

        # 根据return_dict参数决定返回格式
        if not return_dict:
            # 返回元组格式，过滤掉None值
            return tuple(v for v in [hidden_states, encoder_states, all_attentions] if v is not None)
        # 返回结构化的ModelOutput对象
        return BaseModelOutput(
            last_hidden_state=hidden_states,  # 最后一层的隐藏状态
            hidden_states=encoder_states,  # 所有层的隐藏状态
            attentions=all_attentions  # 所有层的注意力权重
        )


class Llama4UnfoldConvolution(nn.Module):
    """Llama4展开卷积模块，用于将图像分割成patch并进行嵌入"""

    def __init__(self, config):
        """初始化展开卷积模块"""
        super().__init__()
        # 获取patch大小（卷积核大小）
        kernel_size = config.patch_size
        # 如果patch_size是整数，转换为(height, width)的元组格式
        if isinstance(kernel_size, int):
            kernel_size = (kernel_size, kernel_size)

        # 创建Unfold层，用于将图像分割成重叠或不重叠的patch
        # kernel_size: patch的大小
        # stride: patch之间的步长，通常等于patch_size以避免重叠
        self.unfold = torch.nn.Unfold(kernel_size=kernel_size, stride=config.patch_size)

        # 线性投影层：将展开后的patch向量映射到隐藏维度
        # 输入维度：通道数 × patch高度 × patch宽度
        # 输出维度：隐藏层大小
        self.linear = nn.Linear(
            config.num_channels * kernel_size[0] * kernel_size[1],  # 展开后每个patch的特征数
            config.hidden_size,  # 目标隐藏维度
            bias=False,  # 不使用偏置
        )

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        """前向传播：图像 -> patch展开 -> 线性投影"""
        # 使用Unfold操作将图像分割成patch
        # 输入：(batch_size, channels, height, width)
        # 输出：(batch_size, channels*patch_h*patch_w, num_patches)
        hidden_states = self.unfold(hidden_states)

        # 调整维度顺序：(batch_size, channels*patch_h*patch_w, num_patches)
        # -> (batch_size, num_patches, channels*patch_h*patch_w)
        hidden_states = hidden_states.permute(0, 2, 1)

        # 通过线性层将每个patch的像素特征映射到隐藏维度
        # 输出：(batch_size, num_patches, hidden_size)
        hidden_states = self.linear(hidden_states)
        return hidden_states


class Llama4VisionRotaryEmbedding(nn.Module):
    """Llama4视觉旋转位置编码模块，为图像patch提供2D位置信息"""

    def __init__(self, config):
        """初始化视觉旋转位置编码"""
        super().__init__()
        # 计算每个维度上的patch数量（假设图像是正方形）
        # 例如：image_size=224, patch_size=16, 则idx=14
        idx = config.image_size // config.patch_size

        # 创建patch索引：将2D patch位置展平为1D索引 (0, 1, 2, ..., idx^2-1)
        # 形状：(idx^2, 1)，每个patch对应一个唯一索引
        img_idx = torch.arange(idx ** 2, dtype=torch.int32).reshape(idx ** 2, 1)

        # 添加一个额外的位置用于CLS token（分类token）
        # 复制第一个索引位置，然后会被替换为特殊标记
        img_idx = torch.cat([img_idx, img_idx[:1]], dim=0)

        # 将最后一个位置标记为CLS token的特殊ID（-2）
        # 这样可以在后续处理中识别并特殊处理CLS token
        img_idx[-1, -1] = -2  # ID_CLS_TOKEN

        # 计算每个patch在2D网格中的x坐标（列索引）
        # 使用模运算获取x坐标：0, 1, 2, ..., idx-1, 0, 1, 2, ...
        frequencies_x = img_idx % idx  # get the coordinates of the 2d matrix along x

        # 计算每个patch在2D网格中的y坐标（行索引）
        # 使用整除获取y坐标：0, 0, 0, ..., 1, 1, 1, ..., 2, 2, 2, ...
        frequencies_y = img_idx // idx  # get the coordinates of the 2d matrix along y

        # 计算每个注意力头中用于RoPE的频率维度数
        # 除以2是因为复数需要成对的cos和sin分量
        freq_dim = config.hidden_size // config.num_attention_heads // 2

        # 计算RoPE的基础频率，使用指数衰减模式
        # rope_theta通常是10000，频率从高到低递减
        rope_freq = 1.0 / (config.rope_theta ** (torch.arange(0, freq_dim, 2)[: (freq_dim // 2)].float() / freq_dim))

        # 为x坐标计算频率矩阵，+1避免零频率
        # repeat_interleave(2)是为了匹配复数表示（实部和虚部）
        freqs_x = ((frequencies_x + 1)[..., None] * rope_freq[None, None, :]).repeat_interleave(2, dim=-1)

        # 为y坐标计算频率矩阵
        freqs_y = ((frequencies_y + 1)[..., None] * rope_freq[None, None, :]).repeat_interleave(2, dim=-1)

        # 将x和y频率连接，形成完整的2D位置编码
        # [::2] 用于提取复数表示中的实部
        freqs = torch.cat([freqs_x, freqs_y], dim=-1).float().contiguous()[..., ::2]

        # 对于特殊token（索引<0，如CLS token），将其频率设为0
        # 这样CLS token就不会受到位置编码的影响
        freqs = freqs.masked_fill(img_idx.reshape(-1, 1, 1) < 0, 0)

        # 将频率转换为复数形式，用于旋转操作
        # 使用欧拉公式：e^(iθ) = cos(θ) + i*sin(θ)
        freq_cis = torch.view_as_complex(torch.stack([torch.cos(freqs), torch.sin(freqs)], dim=-1))

        # 保存频率复数张量，形状说明：(num_patches+1, 1, freq_dim)
        self.freqs_ci = freq_cis  # idx**2, idx**2, idx * 2

    def forward(self, hidden_states):
        """前向传播：返回适配到正确设备的频率复数张量"""
        # 将频率张量移动到与输入hidden_states相同的设备上（CPU/GPU）
        return self.freqs_ci.to(hidden_states.device)


class Llama4VisionModel(Llama4PreTrainedModel):
    """Llama4完整的视觉模型，包含patch embedding、位置编码、编码器等所有组件"""

    # 模型前缀，用于命名和加载
    base_model_prefix = "vision_model"
    # 指定不应该被分割的模块（用于模型并行）
    _no_split_modules = ["Llama4VisionEncoderLayer"]
    # 类型提示：配置对象类型
    config: Llama4VisionConfig

    def __init__(self, config: Llama4VisionConfig):
        """初始化Llama4视觉模型"""
        super().__init__(config)
        # 从配置中提取关键参数
        self.image_size = config.image_size  # 输入图像尺寸
        self.patch_size = config.patch_size  # patch大小
        self.hidden_size = config.hidden_size  # 隐藏层维度
        self.num_channels = config.num_channels  # 图像通道数

        # 计算patch总数：图像patch数 + 1个CLS token
        self.num_patches = (self.image_size // self.patch_size) ** 2 + 1
        # 缩放因子，用于初始化参数，遵循transformer的标准做法
        self.scale = config.hidden_size ** -0.5

        # === 核心组件初始化 ===
        # patch嵌入层：将图像分割成patch并转换为嵌入向量
        self.patch_embedding = Llama4UnfoldConvolution(config)

        # CLS token嵌入：用于分类的特殊token，使用缩放的随机初始化
        self.class_embedding = nn.Parameter(self.scale * torch.randn(self.hidden_size))

        # 位置嵌入：为每个patch（包括CLS token）提供位置信息
        self.positional_embedding_vlm = nn.Parameter(self.scale * torch.randn(self.num_patches, self.hidden_size))

        # 旋转位置编码：提供额外的2D位置信息
        self.rotary_embedding = Llama4VisionRotaryEmbedding(config)

        # === 层归一化组件 ===
        # 编码器前的层归一化
        self.layernorm_pre = nn.LayerNorm(self.hidden_size)
        # 编码器后的层归一化
        self.layernorm_post = nn.LayerNorm(self.hidden_size)

        # === 编码器和适配器 ===
        # 主要的视觉编码器（多层Transformer）
        self.model = Llama4VisionEncoder(config)
        # 视觉适配器：用于最终的特征变换和输出
        self.vision_adapter = Llama4VisionPixelShuffleMLP(config)

        # 执行后初始化（权重初始化等）
        self.post_init()

    def get_input_embeddings(self):
        """
        获取输入嵌入层，用于激活输入的梯度
        这个函数通常用于模型分析和调试
        """
        return self.patch_embedding

    def forward(
            self,
            pixel_values: torch.Tensor,  # 输入的像素值，形状为(batch, channels, height, width)
            attention_mask: Optional[torch.Tensor] = None,  # 注意力掩码（可选）
            output_attentions: Optional[bool] = None,  # 是否输出注意力权重
            output_hidden_states: Optional[bool] = None,  # 是否输出隐藏状态
            return_dict: Optional[bool] = None,  # 返回格式
    ) -> Union[BaseModelOutput, tuple[torch.Tensor, ...]]:
        r"""
        Llama4视觉模型的前向传播

        Example:
        ```python
        >>> from PIL import Image
        >>> import requests
        >>> from transformers import AutoProcessor, MllamaVisionModel

        >>> checkpoint = "meta-llama/Llama-3.2-11B-Vision"
        >>> model = MllamaVisionModel.from_pretrained(checkpoint)
        >>> processor = AutoProcessor.from_pretrained(checkpoint)

        >>> url = "https://www.ilankelman.org/stopsigns/australia.jpg"
        >>> image = Image.open(requests.get(url, stream=True).raw)
        >>> inputs = processor(images=image, return_tensors="pt")

        >>> output = model(**inputs)

        >>> print(output.last_hidden_state.shape)
        torch.Size([1, 1, 4, 1025, 7680])
        ```
        """
        # 设置输出选项的默认值
        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
        output_hidden_states = (
            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
        )
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict

        # === 第一步：Patch Embedding ===
        # 解析输入张量的形状
        batch_size_times_num_tiles, num_channels, height, width = pixel_values.shape
        # 目前这些参数都固定为1，但保留扩展性
        num_concurrent_media = 1  # 并发媒体数量
        num_chunks = 1  # 分块数量

        # 将图像转换为patch嵌入
        # 输出形状：(batch_size_times_num_tiles, num_patches, hidden_dim)
        hidden_state = self.patch_embedding(pixel_values)
        _, num_patches, hidden_dim = hidden_state.shape

        # === 第二步：添加CLS Token ===
        # 重塑张量以准备添加CLS token
        hidden_state = hidden_state.reshape(
            batch_size_times_num_tiles * num_concurrent_media * num_chunks, num_patches, hidden_dim
        )
        # 扩展CLS token以匹配batch size
        class_embedding = self.class_embedding.expand(hidden_state.shape[0], 1, hidden_state.shape[-1])
        # 将CLS token添加到序列的末尾
        hidden_state = torch.cat([hidden_state, class_embedding], dim=1)
        num_patches += 1  # 更新patch总数

        # === 第三步：添加位置编码 ===
        # 重塑张量以适应位置编码的形状
        hidden_state = hidden_state.reshape(
            batch_size_times_num_tiles * num_concurrent_media, num_chunks, num_patches, hidden_dim
        )
        # 获取位置嵌入并转换到正确的数据类型和设备
        positional_embedding = self.positional_embedding_vlm.to(dtype=hidden_state.dtype, device=hidden_state.device)
        # 添加位置编码到隐藏状态
        hidden_state = hidden_state + positional_embedding

        # === 第四步：预处理层归一化 ===
        hidden_state = self.layernorm_pre(hidden_state)

        # === 第五步：准备编码器输入 ===
        # 重塑为编码器期望的形状
        hidden_state = hidden_state.view(batch_size_times_num_tiles, -1, hidden_dim)
        # 获取旋转位置编码的频率
        freqs_ci = self.rotary_embedding(pixel_values)

        # === 第六步：通过编码器 ===
        output = self.model(
            hidden_state,  # 输入隐藏状态
            attention_mask=None,  # 视觉模型通常不使用attention mask
            output_hidden_states=output_hidden_states,  # 是否输出中间隐藏状态
            output_attentions=output_attentions,  # 是否输出注意力权重
            freqs_ci=freqs_ci,  # 旋转位置编码频率
        )

        # === 第七步：后处理 ===
        # 获取编码器的最终输出
        hidden_state = output.last_hidden_state
        # 应用后处理层归一化
        hidden_state = self.layernorm_post(hidden_state)
        # 移除CLS token（保留所有patch的特征）
        hidden_state = hidden_state[:, :-1, :]

        # === 第八步：视觉适配器 ===
        # 使用Pixel Shuffle + MLP进行最终的特征变换
        hidden_state = self.vision_adapter(hidden_state)

        # === 第九步：准备输出 ===
        # 获取所有隐藏状态（如果需要）
        hidden_states = output.hidden_states if output_hidden_states else None

        # 获取注意力权重（如果需要）
        if output_attentions:
            attentions = output[2]
        else:
            attentions = None

        # 根据return_dict决定返回格式
        if not return_dict:
            # 返回元组，过滤掉None值
            return tuple(v for v in [hidden_state, hidden_states, attentions] if v is not None)

        # 返回结构化的BaseModelOutput对象
        return BaseModelOutput(
            last_hidden_state=hidden_state,  # 最终的隐藏状态（去除CLS token后的patch特征）
            hidden_states=hidden_states,  # 所有层的隐藏状态
            attentions=attentions,  # 所有层的注意力权重
        )


class Llama4ForConditionalGeneration(Llama4PreTrainedModel, GenerationMixin):
    """
    Llama4条件生成模型，集成了视觉和语言模型的多模态架构
    继承自预训练模型基类和生成混合类，支持图像-文本的条件生成
    """

    # 指定不应该被分割的模块（用于模型并行）
    _no_split_modules = ["Llama4TextDecoderLayer", "Llama4VisionEncoderLayer"]
    # 张量并行计划（暂时为空）
    _tp_plan = {}
    # 基础模型前缀为空（因为这是顶层模型）
    base_model_prefix = ""
    # 配置类型提示
    config: Llama4Config

    def __init__(self, config: Llama4Config):
        """初始化Llama4条件生成模型"""
        super().__init__(config)

        # === 核心组件初始化 ===
        # 视觉模型：处理图像输入，提取视觉特征
        self.vision_model = Llama4VisionModel(config.vision_config)

        # 多模态投影器：将视觉特征映射到语言模型的特征空间
        self.multi_modal_projector = Llama4MultiModalProjector(config)

        # 语言模型：基于Llama的因果语言模型，负责文本生成
        self.language_model = Llama4ForCausalLM(config.text_config)

        # 词汇表大小，从文本配置中获取
        self.vocab_size = config.text_config.vocab_size

        # 填充token的ID，如果未指定则设为-1
        self.pad_token_id = self.config.pad_token_id if self.config.pad_token_id is not None else -1

        # 执行后初始化处理
        self.post_init()

    def get_input_embeddings(self):
        """获取输入嵌入层，实际上是语言模型的输入嵌入"""
        return self.language_model.get_input_embeddings()

    def set_input_embeddings(self, value):
        """设置输入嵌入层"""
        self.language_model.set_input_embeddings(value)

    def get_output_embeddings(self):
        """获取输出嵌入层（用于生成词汇表logits）"""
        return self.language_model.get_output_embeddings()

    def set_output_embeddings(self, new_embeddings):
        """设置输出嵌入层"""
        self.language_model.set_output_embeddings(new_embeddings)

    def set_decoder(self, decoder):
        """设置解码器（语言模型的解码器部分）"""
        self.language_model.set_decoder(decoder)

    def get_decoder(self):
        """获取解码器"""
        return self.language_model.get_decoder()

    def get_image_features(
            self,
            pixel_values: torch.FloatTensor,  # 输入的像素值
            vision_feature_layer: Union[int, list[int]],  # 要提取特征的层索引
            vision_feature_select_strategy: str,  # 特征选择策略
            **kwargs,
    ):
        """
        从视觉塔获取图像的最后隐藏状态并应用投影

        Args:
            pixel_values (`torch.FloatTensor` of shape `(batch_size, channels, height, width)`):
               对应输入图像的张量
            vision_feature_layer (`Union[int, list[int]]`):
                选择视觉特征的层索引。如果提供多个索引，
                将连接相应索引的视觉特征以形成最终的视觉特征
            vision_feature_select_strategy (`str`):
                从视觉主干网络选择视觉特征的策略。
                可以是 `"default"` 或 `"full"` 之一

        Returns:
            image_features (`torch.Tensor`):
                形状为 `(num_images, image_length, embed_dim)` 的图像特征张量
        """
        # 验证特征选择策略的有效性
        if vision_feature_select_strategy not in ["default", "full"]:
            raise ValueError(f"Unexpected select feature strategy: {self.vision_feature_select_strategy}")

        # 过滤掉None值的kwargs
        kwargs = {k: v for k, v in kwargs.items() if v is not None}

        # 通过视觉模型处理图像，获取图像特征
        image_outputs = self.vision_model(pixel_values, output_hidden_states=False, **kwargs)

        # 提取最后一层的隐藏状态作为图像特征
        hidden_state = image_outputs.last_hidden_state
        return hidden_state

    def get_placeholder_mask(
            self,
            input_ids: torch.LongTensor,  # 输入的token ID序列
            inputs_embeds: torch.FloatTensor,  # 输入的嵌入向量
            image_features: torch.FloatTensor  # 图像特征
    ):
        """
        从 `input_ids` 或 `inputs_embeds` 中获取多模态占位符掩码，
        并检查占位符token数量是否等于多模态特征的长度。
        如果长度不同，将抛出错误。

        占位符掩码用于标识哪些位置需要用图像特征替换文本嵌入
        """
        if input_ids is None:
            # 如果没有input_ids，从inputs_embeds中查找图像token
            # 通过比较嵌入向量是否等于图像token的嵌入来识别
            special_image_mask = inputs_embeds == self.get_input_embeddings()(
                torch.tensor(self.config.image_token_id, dtype=torch.long, device=inputs_embeds.device)
            )
            # 检查是否所有维度都匹配（确保整个嵌入向量都相等）
            special_image_mask = special_image_mask.all(-1)
        else:
            # 如果有input_ids，直接查找图像token ID
            special_image_mask = input_ids == self.config.image_token_id

        # 计算图像token的总数
        n_image_tokens = special_image_mask.sum()

        # 将掩码扩展到与inputs_embeds相同的形状
        special_image_mask = special_image_mask.unsqueeze(-1).expand_as(inputs_embeds).to(inputs_embeds.device)

        # 验证图像特征数量与图像token数量是否匹配
        if inputs_embeds[special_image_mask].numel() != image_features.numel():
            raise ValueError(
                f"Image features and image tokens do not match: tokens: {n_image_tokens}, features {image_features.shape[0]}"
            )
        return special_image_mask

    @auto_docstring
    def forward(
            self,
            input_ids: torch.LongTensor = None,  # 输入token序列
            pixel_values: torch.FloatTensor = None,  # 图像像素值
            attention_mask: Optional[torch.Tensor] = None,  # 注意力掩码
            position_ids: Optional[torch.LongTensor] = None,  # 位置ID
            past_key_values: Optional[Cache] = None,  # 缓存的键值对（用于生成）
            inputs_embeds: Optional[torch.FloatTensor] = None,  # 输入嵌入（可选）
            vision_feature_layer: Optional[Union[int, list[int]]] = None,  # 视觉特征层
            vision_feature_select_strategy: Optional[str] = None,  # 视觉特征选择策略
            labels: Optional[torch.LongTensor] = None,  # 训练标签
            use_cache: Optional[bool] = None,  # 是否使用缓存
            output_attentions: Optional[bool] = None,  # 是否输出注意力权重
            output_hidden_states: Optional[bool] = None,  # 是否输出隐藏状态
            return_dict: Optional[bool] = None,  # 返回字典格式
            cache_position: Optional[torch.LongTensor] = None,  # 缓存位置
            logits_to_keep: Union[int, torch.Tensor] = 0,  # 要保留的logits数量
            image_sizes: torch.Tensor = None,  # 图像尺寸信息
            **kwargs: Unpack[TransformersKwargs],
    ) -> Union[tuple, Llama4CausalLMOutputWithPast]:
        r"""
        多模态条件生成模型的前向传播

        Args:
            labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
                用于计算掩码语言建模损失的标签。索引应该在 `[0, ..., config.vocab_size]`
                或 -100 范围内（参见 `input_ids` 文档字符串）。索引设置为 `-100` 的token
                被忽略（掩码），损失仅对标签在 `[0, ..., config.vocab_size]` 的token计算。

        Example:
        ```python
        >>> from PIL import Image
        >>> import requests
        >>> from transformers import AutoProcessor, LlavaForConditionalGeneration

        >>> model = LlavaForConditionalGeneration.from_pretrained("llava-hf/llava-1.5-7b-hf")
        >>> processor = AutoProcessor.from_pretrained("llava-hf/llava-1.5-7b-hf")

        >>> prompt = "USER: <image>\nWhat's the content of the image? ASSISTANT:"
        >>> url = "https://www.ilankelman.org/stopsigns/australia.jpg"
        >>> image = Image.open(requests.get(url, stream=True).raw)

        >>> inputs = processor(images=image, text=prompt, return_tensors="pt")

        >>> # 生成
        >>> generate_ids = model.generate(**inputs, max_new_tokens=15)
        >>> processor.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
        "USER:  \nWhat's the content of the image? ASSISTANT: The image features a busy city street with a stop sign prominently displayed"
        ```
        """

        # === 第一步：设置默认参数 ===
        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
        output_hidden_states = (
            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
        )
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict

        # 设置视觉特征相关的默认参数
        vision_feature_layer = (
            vision_feature_layer
            if vision_feature_layer is not None
            else self.config.vision_config.vision_feature_layer
        )
        vision_feature_select_strategy = (
            vision_feature_select_strategy
            if vision_feature_select_strategy is not None
            else self.config.vision_config.vision_feature_select_strategy
        )

        # === 第二步：输入验证 ===
        # 确保input_ids和inputs_embeds只指定其中一个
        if (input_ids is None) ^ (inputs_embeds is not None):
            raise ValueError("You must specify exactly one of input_ids or inputs_embeds")

        # 确保pixel_values和inputs_embeds不能同时指定
        if pixel_values is not None and inputs_embeds is not None:
            raise ValueError(
                "You cannot specify both pixel_values and inputs_embeds at the same time, and must specify either one"
            )

        # === 第三步：准备输入嵌入 ===
        # 如果没有提供inputs_embeds，从input_ids生成
        if inputs_embeds is None:
            inputs_embeds = self.get_input_embeddings()(input_ids)

        # === 第四步：处理图像输入（如果存在） ===
        if pixel_values is not None:
            # 提取图像特征
            image_features = self.get_image_features(
                pixel_values=pixel_values,
                vision_feature_layer=vision_feature_layer,
                vision_feature_select_strategy=vision_feature_select_strategy,
                image_sizes=image_sizes,
            )

            # 将图像特征展平并通过多模态投影器
            vision_flat = image_features.view(-1, image_features.size(-1))
            projected_vision_flat = self.multi_modal_projector(vision_flat).to(
                inputs_embeds.device, inputs_embeds.dtype
            )

            # 获取占位符掩码，确定哪些位置需要替换为图像特征
            special_image_mask = self.get_placeholder_mask(
                input_ids, inputs_embeds=inputs_embeds, image_features=projected_vision_flat
            )

            # 使用图像特征替换对应的文本嵌入
            inputs_embeds = inputs_embeds.masked_scatter(special_image_mask, projected_vision_flat)

        # === 第五步：通过语言模型进行前向传播 ===
        outputs = self.language_model(
            attention_mask=attention_mask,  # 注意力掩码
            position_ids=position_ids,  # 位置ID
            past_key_values=past_key_values,  # 缓存的键值对
            inputs_embeds=inputs_embeds,  # 融合了图像信息的输入嵌入
            use_cache=use_cache,  # 是否使用缓存
            output_attentions=output_attentions,  # 是否输出注意力权重
            output_hidden_states=output_hidden_states,  # 是否输出隐藏状态
            return_dict=return_dict,  # 返回格式
            cache_position=cache_position,  # 缓存位置
            logits_to_keep=logits_to_keep,  # 要保留的logits
            **kwargs,
        )

        # 获取输出logits
        logits = outputs[0]

        # === 第六步：计算损失（如果提供了标签） ===
        loss = None
        if labels is not None:
            # 移位操作：让位置 < n 的token预测位置 n 的token
            if attention_mask is not None:
                # 使用输入注意力掩码来移位logits和labels，因为它是2D的
                # 也会裁剪注意力掩码，以防它更长（在PrefixTuning中会发生）
                shift_attention_mask = attention_mask[:, -(logits.shape[1] - 1):].to(logits.device)
                shift_logits = logits[..., :-1, :][shift_attention_mask.to(logits.device) != 0].contiguous()
                shift_labels = labels[..., 1:][shift_attention_mask.to(labels.device) != 0].contiguous()
            else:
                # 如果没有注意力掩码，简单地移位
                shift_logits = logits[..., :-1, :].contiguous()
                shift_labels = labels[..., 1:].contiguous()

            # 展平token并计算交叉熵损失
            loss_fct = nn.CrossEntropyLoss()
            loss = loss_fct(
                shift_logits.view(-1, shift_logits.size(-1)),
                shift_labels.view(-1).to(shift_logits.device)
            )

        # === 第七步：准备输出 ===
        if not return_dict:
            # 返回元组格式
            output = (logits,) + outputs[1:]
            return (loss,) + output if loss is not None else output

        # 返回结构化输出
        return Llama4CausalLMOutputWithPast(
            loss=loss,  # 计算的损失
            logits=logits,  # 输出logits
            past_key_values=outputs.past_key_values,  # 缓存的键值对
            hidden_states=outputs.hidden_states,  # 隐藏状态
            attentions=outputs.attentions,  # 注意力权重
            image_hidden_states=image_features if pixel_values is not None else None,  # 图像隐藏状态
        )

    def prepare_inputs_for_generation(
            self,
            input_ids,  # 输入token ID
            past_key_values=None,  # 过去的键值对缓存
            inputs_embeds=None,  # 输入嵌入
            pixel_values=None,  # 像素值
            attention_mask=None,  # 注意力掩码
            cache_position=None,  # 缓存位置
            logits_to_keep=None,  # 要保留的logits
            **kwargs,
    ):
        """
        为生成准备输入 - 重写方法
        在特定情况下，我们不希望将图像输入转发给模型

        这个方法处理缓存解码阶段的特殊逻辑：
        - 在缓存解码阶段，图像已经被处理过，不需要重复处理
        - 只有在生成的第一步（cache_position[0] == 0）才需要图像输入
        """

        # 调用语言模型的输入准备方法
        model_inputs = self.language_model.prepare_inputs_for_generation(
            input_ids,
            past_key_values=past_key_values,
            inputs_embeds=inputs_embeds,
            attention_mask=attention_mask,
            cache_position=cache_position,
            logits_to_keep=logits_to_keep,
            **kwargs,
        )

        # 只有在生成的第一步才传递像素值
        if cache_position[0] == 0:
            # 如果我们在缓存解码阶段，像素值应该是None，
            # 因为input_ids不再包含特殊图像token
            # 否则我们需要将像素值传递给模型
            model_inputs["pixel_values"] = pixel_values

        return model_inputs


__all__ = [
    "Llama4PreTrainedModel",
    "Llama4TextModel",
    "Llama4VisionModel",
    "Llama4ForCausalLM",
    "Llama4ForConditionalGeneration",
]

